from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation

gpu_list = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal

def render(grid, state, goal, valid = True):
        
        state_x = int(state[0])
        state_y = int(state[1])

        food_x = int(goal[0])
        food_y = int(goal[1])

        grid_size_display = 20
        width, height = grid.shape[0], grid.shape[1]
        img = np.zeros((width * grid_size_display, height * grid_size_display, 3), np.uint8)
        
        for j in range(width):
            for i in range(height):
                if grid[j,i] == 1:
                    cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), (255, 255, 255), -1)
                    # draw border with color(100,100,100)
                    cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), (100, 100, 100), 1)
                else:
                    cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), (0, 0, 0), -1)
                    # draw border with color(100,100,100)
                    cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), (100, 100, 100), 1)
                if j == state_x and i == state_y:
                    cv2.circle(img, (i * grid_size_display + int(grid_size_display/2), j * grid_size_display + int(grid_size_display/2)), 7, (0, 0, 255), -1, cv2.LINE_AA)
        
        # put with a dot on food position
        cv2.circle(img, (food_y * grid_size_display + grid_size_display//2, food_x * grid_size_display + grid_size_display//2), 7, (0,100,0), -1, cv2.LINE_AA)

        # put with a dot on food position
        cv2.circle(img, (food_y * grid_size_display + grid_size_display//2, food_x * grid_size_display + grid_size_display//2), 7, (0,100,0), -1, cv2.LINE_AA)

        if not valid:
            # draw a big red cross
            cv2.line(img, (0, 0), (img.shape[1], img.shape[0]), (0, 0, 255), 5, cv2.LINE_AA)
            cv2.line(img, (0, img.shape[0]), (img.shape[1], 0), (0, 0, 255), 5, cv2.LINE_AA)
            cv2.putText(img, "invalid map", (int(img.shape[1]/2) - 100, int(img.shape[0]/2)), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (200, 200, 200), 3, cv2.LINE_AA)
            cv2.putText(img, "invalid map", (int(img.shape[1]/2) - 100, int(img.shape[0]/2)), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (200, 0, 0), 2, cv2.LINE_AA)
        return img

event_type = ""
event_x = 0
event_y = 0

def run_editor(landscape, state, goal, map_size=12):

    global event_type, event_x, event_y

    # define mouse callback function
    def input_cb(event,x,y,flags,param):
        global event_type, event_x, event_y
        if event == cv2.EVENT_MBUTTONUP:
            event_x = x
            event_y = y
            event_type = "flip_space"
        elif event == cv2.EVENT_LBUTTONUP:
            event_x = x
            event_y = y
            event_type = "set_start"
        elif event == cv2.EVENT_RBUTTONUP:
            event_x = x
            event_y = y
            event_type = "set_goal"

    cv2.namedWindow("img", cv2.WINDOW_GUI_NORMAL)
    cv2.setMouseCallback("img", input_cb)

    grid = np.array(landscape).reshape(map_size, map_size).transpose()
    valid = True
    grid_size_display = 20

    while True:

        if event_type == "flip_space":
            grid[event_y//grid_size_display, event_x//grid_size_display] = 1 - grid[event_y//grid_size_display, event_x//grid_size_display]
            event_type = "flip_space_done"

        num_labels, labels, stats, centroids, num_freespace, landscape_img = check_num_labels(grid, map_size, map_size)
        non_zeros = np.count_nonzero(landscape)
        if not (num_labels == 2 and non_zeros >= 5):
            valid = False
        else:
            valid = True
        
        img = render(grid, state, goal, valid)
        cv2.imshow("img", img)
        k = cv2.waitKey(1)
        if k == ord('q'):
            break
        elif k == ord('r'):
            grid = np.array(landscape).reshape(map_size, map_size).transpose()
    
    grid_ = grid.transpose()
    grid_ = grid_.reshape(map_size*map_size).tolist()
    return grid_

def qc_single_goal_record(record):

    var_threshold = 10
    optimization_scale = 2

    var_threshold = 5
    optimization_scale = 2

    trial_len = 0
    trial_len_old = 0
    step_ = 0
    trial_len_var = 0
    first_trial_len = -1
    for i in range(record.shape[0]):
        step_ += 1
        if record[i] == True:
            if first_trial_len == -1:
                first_trial_len = step_
            trial_len = step_
            step_ = 0
            # compute the variance of trial_len
            trial_len_var = 0.4 * abs(trial_len - trial_len_old) + 0.6 * trial_len_var
            trial_len_old = trial_len
            if trial_len_var <= var_threshold and first_trial_len/trial_len >= optimization_scale:
                return True
    return False

def main():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--device", type=int, default=rpl_config.device)
    parser.add_argument("--probe_point", type=int, default=rpl_config.probe_point)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.device = args.device
    rpl_config.probe_point = args.probe_point

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)
    
    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 5000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)
    rnn_state_old = rnn_state.copy()
    diff = jnp.abs(rnn_state - rnn_state_old)
    rnn_state_old = rnn_state.copy()
    diff_norm = jnp.linalg.norm(diff, axis=1)
    diff_norm_old = diff_norm.copy()
    norm_std = diff_norm.copy()

    rnn_state_init = rnn_state.copy()

    obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

    rnn_state_trajectory = []

    for t in range(rpl_config.life_duration):

        if t == rpl_config.probe_point:
            rnn_state_init = rnn_state.copy()

        progress_bar(t, rpl_config.life_duration)

        """ model forward 
        """
        rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
        diff = jnp.abs(rnn_state - rnn_state_old)
        rnn_state_old = rnn_state.copy()
        diff_norm = jnp.linalg.norm(diff, axis=1)
        norm_std = 0.4 * jnp.abs(diff_norm - diff_norm_old) + 0.6 * norm_std
        diff_norm_old = diff_norm.copy()

        rnn_state_trajectory.append(np.array(rnn_state).copy())
            
    print(rnn_state.shape)
    print(norm_std.shape)
    rnn_state_np = np.array(rnn_state)

    rnn_state_trajectory = np.array(rnn_state_trajectory)
    print("rnn_state_trajectory shape: ", rnn_state_trajectory.shape)

    # 显示 rnn_state 的内存大小
    print("rnn_state_trajectory mem：{} MB".format(rnn_state_trajectory.nbytes / 1024 / 1024))

    # 计算 rnn_state 的第二个维度（即每个通道）的方差
    variances = np.var(rnn_state_np, axis=0)
    print(variances)

    plt.plot(norm_std)
    plt.title('Var')
    plt.show()

    # # 从 rnn_state_np 中随机选择 4 个向量
    # random_vectors = random.sample(list(rnn_state_np), 4)
    # # 创建一个包含 4 个子图的图形
    # fig, axs = plt.subplots(2, 2, figsize=(10, 10))

    # # 在每个子图中绘制柱状图
    # for i, vector in enumerate(random_vectors):
    #     ax = axs[i // 2][i % 2]
    #     ax.bar(range(len(vector)), vector)
    #     ax.set_title(f'Vector {i+1}')

    # # 显示图形
    # plt.show()

    # 对 rnn_state_np 进行 PCA
    pca = PCA()
    pca.fit(rnn_state_np)
    rnn_state_pca = pca.transform(rnn_state_np)

    for i in range(rnn_state_trajectory.shape[0]):
        progress_bar(i, rnn_state_trajectory.shape[0])
        rnn_state_trajectory[i] = pca.transform(rnn_state_trajectory[i])

    print("shape of rnn_state_trajectory_pca: ", rnn_state_trajectory.shape)

    # 打印每个主成分的方差贡献率和累计方差贡献率
    variance_ratios = pca.explained_variance_ratio_
    cumulative_variances = np.cumsum(variance_ratios)
    for i, (variance_ratio, cumulative_variance) in enumerate(zip(variance_ratios, cumulative_variances)):
        print(f'PC{i+1} variance ratio: {variance_ratio:.4f}, cumulative variance: {cumulative_variance:.4f}')

    # 创建一个 3D 图形
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # # 绘制前三个主成分的值为散点图
    # ax.scatter(rnn_state_pca[:, 0], rnn_state_pca[:, 1], rnn_state_pca[:, 2])

    # # 设置坐标轴标签
    # ax.set_xlabel('PC1')
    # ax.set_ylabel('PC2')
    # ax.set_zlabel('PC3')

    # # 显示图形
    # plt.show()

    # 创建一个空的散点图
    scatter = ax.scatter([], [], [])

    # 设置坐标轴标签
    ax.set_xlabel('PC1')
    ax.set_ylabel('PC2')
    ax.set_zlabel('PC3')

    view_dim0 = 6
    view_dim1 = 7
    view_dim2 = 8

    # 定义更新函数
    def update(frame):
        progress_bar(frame, rnn_state_trajectory.shape[0])
        # 更新散点图的数据
        scatter._offsets3d = (rnn_state_trajectory[frame, :, view_dim0], rnn_state_trajectory[frame, :, view_dim1], rnn_state_trajectory[frame, :, view_dim2])
        return scatter,

    # 将 rnn_state_trajectory 裁剪到目前尺寸的1/3
    rnn_state_trajectory = rnn_state_trajectory[:rnn_state_trajectory.shape[0] // 3]

    # 创建动画
    ani = FuncAnimation(fig, update, frames=rnn_state_trajectory.shape[0], interval=1, blit=False)

    # 显示动画
    plt.show()

    return
    
    # 对 rnn_state_np 进行降维
    tsne = TSNE(n_components=2, random_state=0)
    rnn_state_tsne = tsne.fit_transform(rnn_state_np)

    rnn_state_init_tsne = tsne.fit_transform(rnn_state_init)

    # 创建两个子图
    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(18, 6))

    # 在第一个子图中绘制 rnn_state_init_tsne
    scatter1 = ax1.scatter(rnn_state_tsne[:, 0], rnn_state_tsne[:, 1])
    ax2.set_title('rnn_states ' + str(rpl_config.life_duration - rpl_config.probe_point) + ' steps ealier')

    # 在第二个子图中绘制 rnn_state_tsne
    scatter2 = ax2.scatter(rnn_state_init_tsne[:, 0], rnn_state_init_tsne[:, 1])
    ax1.set_title('rnn_states on ' + str(rpl_config.life_duration) + 'th step')

    edge_colors = ['blue'] * rnn_state_init_tsne.shape[0]

    contraction_rad = 6

    # 获取 rnn_state_tsne 在 x/y 方向上的最小值和最大值
    x_min, x_max = np.amin(rnn_state_tsne[:, 0]), np.amax(rnn_state_tsne[:, 0])
    y_min, y_max = np.amin(rnn_state_tsne[:, 1]), np.amax(rnn_state_tsne[:, 1])
    max_distances = np.zeros((int(y_max) - int(y_min), int(x_max) - int(x_min)))
    # 将 max_distances 中的值初始化为 100
    # max_distances.fill(100)

    for y in range(int(y_min), int(y_max)):
        for x in range(int(x_min), int(x_max)):
            progress_bar((y-int(y_min)) * (int(x_max) - int(x_min)) + x-int(x_min), (int(x_max) - int(x_min)) * (int(y_max) - int(y_min)))

            # 计算距离最近的数据点
            distances = np.sqrt((x - scatter1.get_offsets()[:, 0]) ** 2 + (y - scatter1.get_offsets()[:, 1]) ** 2)
            indices = np.where(distances < contraction_rad)[0]

            if len(indices)>0:
                # 高亮轨迹源点
                rnn_state_init_source = rnn_state_init_tsne[indices]
                # rnn_state_init_source = rnn_state_init[indices]

                # rnn_state_init_source 是包含点坐标的集合
                points = np.array(rnn_state_init_source)
                # 计算每个点之间的距离
                distances = np.sqrt(((points[:, np.newaxis] - points) ** 2).sum(axis=2))
                # 找到距离最远的两个点
                i, j = np.unravel_index(np.argmax(distances), distances.shape)
                # 计算这两个点之间的距离
                max_distance = distances[i, j]
                # 将 max_distance 值存储在相应的数组元素中
                max_distances[y - int(y_min), x - int(x_min)] = max_distance

                # variances = np.var(rnn_state_init_source, axis=0)
                # variance = np.mean(variances)
                # max_distances[y - int(y_min), x - int(x_min)] = variance

    # contraction map
    max_distances = np.max(max_distances) - max_distances
    im = ax3.imshow(max_distances, cmap='RdBu', origin='lower', extent=[x_min, x_max, y_min, y_max])
    plt.colorbar(im, ax=ax3)

    # 定义鼠标悬停事件处理函数
    def on_hover(event):
        if event.inaxes == ax1:
            # 删除以前绘制的圆圈
            for child in ax1.get_children():
                if isinstance(child, matplotlib.patches.Circle):
                    child.remove()
            for child in ax2.get_children():
                if isinstance(child, matplotlib.patches.Circle):
                    child.remove()
            # 获取鼠标位置
            x, y = event.xdata, event.ydata
            # 计算距离最近的数据点
            distances = np.sqrt((x - scatter1.get_offsets()[:, 0]) ** 2 + (y - scatter1.get_offsets()[:, 1]) ** 2)
            index = np.argmin(distances)
            indices = np.where(distances < contraction_rad)[0]

            rnn_state_init_source = rnn_state_init_tsne[indices]

            # 高亮轨迹源点
            for i in range(len(indices)):
                circle = matplotlib.patches.Circle((rnn_state_init_source[i,0], rnn_state_init_source[i,1]), radius=2, color='red', fill=True)
                ax2.add_patch(circle)
            
            # 绘制圆圈
            circle = matplotlib.patches.Circle((x, y), radius=contraction_rad, color='red', fill=False, linewidth=2)
            ax1.add_patch(circle)

            fig.canvas.draw_idle()

    # 绑定鼠标悬停事件处理函数
    fig.canvas.mpl_connect('motion_notify_event', on_hover)

    # 显示 plot
    plt.show()

            
if __name__ == "__main__":
    main()